from Goldenberry.optimization.cost_functions import *
from Goldenberry.optimization.base.GbCostFunction import GbCostFunction

import numpy as np
from unittest import *

class CostFunctionTest(TestCase):

    def test_onemax(self):
        sample = np.array([[0,1,0,1, 0]])
        func = GbCostFunction(Binary_OneMax)
        self.assertEqual(func(sample), 2)

    def test_statistics(self):
        sample1, cost1 = np.array([[0,1,1,1,0]]), 3.0
        sample2, cost2 = np.array([[0,0,0,1,1]]), 2.0
        sample3, cost3 = np.array([[0,1,1,1,1]]), 4.0
        func = GbCostFunction(Binary_OneMax)
        exp_cost1 = func(sample1)
        exp_cost2 = func(sample2)
        exp_cost3 = func(sample3)
        evals, argmin, argmax, min, max, mean, stdev = func.statistics()
        self.assertEqual((evals, argmin, argmax, min, max, mean), (3,2,3,cost2,cost3,3.0))
        self.assertAlmostEqual(stdev, np.sqrt(2.0/3.0), 5)

    def test_zero(self):
        sample = np.array([[0,1,0,1,0]])
        func = GbCostFunction(Binary_ZeroMax)
        self.assertEqual(func(sample), 3)

    def test_LeadingOnesBlocks(self):
        sample = np.array([1, 1, 1, 0, 0, 0, 1, 0, 1])
        self.assertEqual(Binary_LeadingOnesBlocks(sample), 1)
        sample = np.array([1, 1, 1, 1, 0, 1, 1, 0, 1])
        self.assertEqual(Binary_LeadingOnesBlocks(sample), 1)
        sample = np.array([1, 1, 1, 1, 1, 1, 0, 1, 0])
        self.assertEqual(Binary_LeadingOnesBlocks(sample), 2)
        sample = np.array([1, 1, 1, 1, 1, 1, 1, 1, 1])
        self.assertEqual(Binary_LeadingOnesBlocks(sample), 3)
    
    #def test_condonemax_base(self):
    #    length = 10
    #    func = CondOnemax()
    #    sample = np.array([[0,1,1,1,1,0,1,0,1,1]])
    #    self.assertEqual(func.cost(sample), 4.0)

    #def test_condonemax_base_1(self):
    #    func = CondOnemax(cond_indexes =[9,8,7,6,5,4,3,2,1,0])
    #    self.assertTrue(np.equal(func.cond_indexes, [9,8,7,6,5,4,3,2,1,0]).all())

    #def test_condonemax_cost_calculation(self):
    #    sample = np.array([[0,1,1,1,1,0,1,0,1,1]])
    #    func = CondOnemax(cond_indexes = [9,8,7,6,5,4,3,2,1,0])
    #    self.assertEqual(func.cost(sample)[0], 5.0)
    
    def test_custom_function(self):
        sample = np.array([[0,1,0,1,0]])
        cust = GbCostFunction(script = "def mycustomfunction(solution):\n\treturn solution.sum() - 1.0")
        self.assertEqual(cust.cost(sample), 1.0)

    def test_rastrigin(self):
        #sample = np.array([[0.2,1.5,5.12,3.4,-5.12]])
        sample = np.array([[0.2,1.5,5.12,3.4,-5.12]])
        func = GbCostFunction(Real_Rastrigin)
        self.assertEqual(func(sample), -116.69942745157178)
    
    def test_sphere(self):
        sample = np.array([[0.2,1.5,5.12,3.4,-5.12]])
        func = GbCostFunction(Real_Sphere)
        self.assertEqual(func(sample), -66.2788)

    def test_step(self):
        sample = np.array([[0.2,1.5,5.12,3.4,-5.12]])
        func = GbCostFunction(Real_Step)
        self.assertEqual(func(sample), -45.34)

    def test_schwefel(self):
        sample = np.array([[0.2,-512.3,511.97,3.4,200.1]])
        func = GbCostFunction(Real_Schwefel)
        self.assertEqual(round(func(sample),8),206.64681438)

    def test_griewank(self):
        sample = np.array([[-600.0,600.0,102.5,3.4,0]])
        func = GbCostFunction(Real_Griewank)
        self.assertEqual(round(func(sample),8), -183.51847534)

if __name__ == '__main__':
    unittest.main()
